package Solutions

type cou_uf struct {
	parent []int
	rank   []int
}

func (c *cou_uf) find(i int) int {
	if c.parent[i] != i {
		t := c.find(c.parent[i])
		c.parent[i] = t
		return t
	}
	return c.parent[i]
}
func (c *cou_uf) union(i, j int) {
	i1, j1 := c.find(i), c.find(j)
	ri, rj := c.rank[i1], c.rank[j1]
	if ri == rj {
		c.parent[i1] = j1
		c.rank[j1]++
	} else if ri < rj {
		c.parent[i1] = j1
	} else {
		c.parent[j1] = i1
	}
}

func countPairs__(n int, edges [][]int) int64 {
	if len(edges) == 0 {
		return int64(n*(n-1)) / 2
	}
	var uf = cou_uf{}
	uf.parent = make([]int, n)
	uf.rank = make([]int, n)
	for i := 0; i < n; i++ {
		uf.parent[i] = i
	}
	for i := 0; i < len(edges); i++ {
		uf.union(edges[i][0], edges[i][1])
	}
	var mp = make(map[int]int)
	for i := 0; i < n; i++ {
		mp[uf.find(i)]++
	}
	ans := 0
	if len(mp) == 1 {
		return 0
	}
	for _, v := range mp {
		ans += (n - v) * v
	}
	return int64(ans / 2)
}
